package com.aiwiown.snackmq.common.util;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.protocol.pb.SnackMessageProto;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import lombok.extern.slf4j.Slf4j;

import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 一个用于在领域模型 Message 对象和 Protobuf DTO (SnackMessageProto.SnackMessage) 之间进行转换的工具类。
 */
@Slf4j
public final class MessageConverter {

    private static final AtomicInteger poolSize = new AtomicInteger(0);
    private static final int MAX_POOL_SIZE = 1000;

    private MessageConverter() {
    }

    public static SnackMessageProto.SnackMessage toProto(Message domainMessage) {
        if (domainMessage == null) {
            return SnackMessageProto.SnackMessage.getDefaultInstance();
        }
        SnackMessageProto.SnackMessage.Header.Builder headerBuilder = SnackMessageProto.SnackMessage.Header.newBuilder()
                .setMessageId(String.valueOf(domainMessage.getMessageId()))
                .setTimestamp(domainMessage.getCreateTime());

        if (domainMessage.getTopic() != null) headerBuilder.setTopic(domainMessage.getTopic());
        if (domainMessage.getKey() != null) headerBuilder.setKey(domainMessage.getKey());
        if (domainMessage.getTag() != null) headerBuilder.setTag(domainMessage.getTag());
        if (domainMessage.getErrorMessage() != null) headerBuilder.setErrorMessage(domainMessage.getErrorMessage());

        // 【最终修复】: 增强属性序列化的健壮性，防止因 Map 中存在 null 值导致 NullPointerException。
        // Protobuf 的 Map 不允许 null key 或 null value，这里增加过滤。
        if (domainMessage.getProperties() != null) {
            domainMessage.getProperties().entrySet().stream()
                    .filter(entry -> entry.getKey() != null && entry.getValue() != null)
                    .forEach(entry -> headerBuilder.putProperties(entry.getKey(), entry.getValue()));
        }


        if (domainMessage.getSource() != null) headerBuilder.setSource(domainMessage.getSource());
        if (domainMessage.getTarget() != null) headerBuilder.setTarget(domainMessage.getTarget());

        if (domainMessage.getType() != null) {
            headerBuilder.setType(toProto(domainMessage.getType()));
        }
        if (domainMessage.getStatus() != null) {
            headerBuilder.setStatus(toProto(domainMessage.getStatus()));
        }
        SnackMessageProto.SnackMessage.Builder protoBuilder = SnackMessageProto.SnackMessage.newBuilder()
                .setHeader(headerBuilder);
        if (domainMessage.getBody() != null) {
            protoBuilder.setBody(ByteString.copyFrom(domainMessage.getBody()));
        }

        if (domainMessage.isBatch()) {
            domainMessage.getMessages().stream()
                    .map(MessageConverter::toProto)
                    .forEach(protoBuilder::addMessages);
        }

        return protoBuilder.build();
    }

    /**
     * 【新增】一个重载的 toProto 方法，允许在序列化时注入额外的属性。
     * 这是避免在状态机中创建大量新 Message 对象的关键。
     *
     * @param domainMessage   原始的领域模型消息
     * @param extraProperties 要合并到消息头中的额外属性
     * @return 构建好的 Protobuf 消息
     */
    public static SnackMessageProto.SnackMessage toProto(Message domainMessage, Map<String, String> extraProperties) {
        if (domainMessage == null) {
            return SnackMessageProto.SnackMessage.getDefaultInstance();
        }
        // --- 【最终性能优化】: 直接构建，避免 "toProto().toBuilder()" 带来的巨大开销 ---
        // 这是消除状态机中 CPU 瓶颈的最终、也是最关键的一步。
        // 旧的实现会先创建一个完整的 ProtoBuf 对象，然后再将其转回 Builder，这在单线程关键路径上是不可接受的。
        // 新的实现直接从原始 Message 对象构建，只在最后调用一次 build()。
        SnackMessageProto.SnackMessage.Header.Builder headerBuilder = SnackMessageProto.SnackMessage.Header.newBuilder()
                .setMessageId(String.valueOf(domainMessage.getMessageId()))
                .setTimestamp(domainMessage.getCreateTime());
        if (domainMessage.getTopic() != null) headerBuilder.setTopic(domainMessage.getTopic());
        if (domainMessage.getKey() != null) headerBuilder.setKey(domainMessage.getKey());
        if (domainMessage.getTag() != null) headerBuilder.setTag(domainMessage.getTag());
        if (domainMessage.getErrorMessage() != null) headerBuilder.setErrorMessage(domainMessage.getErrorMessage());
        if (domainMessage.getSource() != null) headerBuilder.setSource(domainMessage.getSource());
        if (domainMessage.getTarget() != null) headerBuilder.setTarget(domainMessage.getTarget());
        if (domainMessage.getType() != null) headerBuilder.setType(toProto(domainMessage.getType()));
        if (domainMessage.getStatus() != null) headerBuilder.setStatus(toProto(domainMessage.getStatus()));
        // 先合并原始属性
        if (domainMessage.getProperties() != null) {
            headerBuilder.putAllProperties(domainMessage.getProperties());
        }
        // 再合并额外属性（如果 key 相同，会覆盖原始属性，这符合我们的预期）
        if (extraProperties != null && !extraProperties.isEmpty()) {
            headerBuilder.putAllProperties(extraProperties);
        }
        if (domainMessage.getStatus() != null) headerBuilder.setStatus(toProto(domainMessage.getStatus()));

        // 【最终修复】: 同样增强属性序列化的健壮性。
        // 先合并原始属性
        if (domainMessage.getProperties() != null) {
            domainMessage.getProperties().entrySet().stream()
                    .filter(entry -> entry.getKey() != null && entry.getValue() != null)
                    .forEach(entry -> headerBuilder.putProperties(entry.getKey(), entry.getValue()));
        }
        // 再合并额外属性（如果 key 相同，会覆盖原始属性，这符合我们的预期）
        if (extraProperties != null && !extraProperties.isEmpty()) {
            extraProperties.entrySet().stream()
                    .filter(entry -> entry.getKey() != null && entry.getValue() != null)
                    .forEach(entry -> headerBuilder.putProperties(entry.getKey(), entry.getValue()));
        }
        SnackMessageProto.SnackMessage.Builder protoBuilder = SnackMessageProto.SnackMessage.newBuilder().setHeader(headerBuilder);
        if (domainMessage.getBody() != null) {
            protoBuilder.setBody(ByteString.copyFrom(domainMessage.getBody()));
        }
        // 递归处理批量消息
        if (domainMessage.isBatch()) {
            domainMessage.getMessages().stream()
                    .map(MessageConverter::toProto) // 注意：这里调用的是无额外属性的基础方法
                    .forEach(protoBuilder::addMessages);
        }
        return protoBuilder.build();
    }

    public static Message fromProto(SnackMessageProto.SnackMessage protoMessage) {
        if (protoMessage == null) {
            return null;
        }

        SnackMessageProto.SnackMessage.Header header = protoMessage.getHeader();
        Message.MessageBuilder domainBuilder = Message.builder()
                .messageId(header.getMessageId() != null && !header.getMessageId().isEmpty() ? Long.parseLong(header.getMessageId()) : 0)
                .createTime(header.getTimestamp())
                .topic(header.getTopic())
                .key(header.getKey())
                .tag(header.getTag())
                .errorMessage(header.getErrorMessage())
                // 【核心修复】: 将 Protobuf 的不可变 Map 转换为可变的 HashMap。
                // 这可以防止在 ConsumeMessageHandler 中调用 addProperty 时抛出 UnsupportedOperationException。
                .properties(new HashMap<>(header.getPropertiesMap()))
                .source(header.getSource())
                .target(header.getTarget());

        domainBuilder.type(fromProto(header.getType()));
        domainBuilder.status(fromProto(header.getStatus()));

        domainBuilder.body(protoMessage.getBody().toByteArray());

        if (protoMessage.getMessagesCount() > 0) {
            domainBuilder.messages(
                    protoMessage.getMessagesList().stream()
                            .map(MessageConverter::fromProto)
                            .collect(Collectors.toList())
            );
        }

        return domainBuilder.build();
    }


    public static byte[] toBytes(Message message) {
        if (message == null) {
            return null;
        }
        try {
            SnackMessageProto.SnackMessage protoMessage = toProto(message);
            return protoMessage.toByteArray();
        } catch (Exception e) {
            log.error("Failed to convert Message to bytes", e);
            return null;
        }
    }


    /**
     * 【新增】一个重载的 toBytes 方法，允许在序列化时注入额外的属性。
     *
     * @param message         原始消息
     * @param extraProperties 额外属性
     * @return 序列化后的字节数组
     */
    public static byte[] toBytes(Message message, Map<String, String> extraProperties) {
        if (message == null) {
            return null;
        }
        try {
            SnackMessageProto.SnackMessage protoMessage = toProto(message, extraProperties);
            return protoMessage.toByteArray();
        } catch (Exception e) {
            log.error("Failed to convert Message with extra properties to bytes", e);
            return null;
        }
    }

    public static Message fromBytes(byte[] bytes) {
        if (bytes == null || bytes.length == 0) {
            return null;
        }
        try {
            SnackMessageProto.SnackMessage protoMessage = SnackMessageProto.SnackMessage.parseFrom(bytes);
            return fromProto(protoMessage);
        } catch (Exception e) {
            log.error("Failed to convert bytes to Message", e);
            return null;
        }
    }

    public static Message fromByteBuffer(ByteBuffer buffer) {
        if (buffer == null || !buffer.hasRemaining()) {
            return null;
        }
        try {
            SnackMessageProto.SnackMessage protoMessage = SnackMessageProto.SnackMessage.parseFrom(buffer);
            return fromProto(protoMessage);
        } catch (InvalidProtocolBufferException e) {
            log.error("Failed to convert ByteBuffer to Message due to invalid protobuf data", e);
            return null;
        } catch (Exception e) {
            log.error("An unexpected error occurred while converting ByteBuffer to Message", e);
            return null;
        }
    }

    private static SnackMessageProto.ProtoMessageType toProto(MessageType type) {
        if (type == null) return SnackMessageProto.ProtoMessageType.UNKNOWN;
        switch (type) {
            case PRODUCE_REQUEST:
                return SnackMessageProto.ProtoMessageType.PRODUCE_REQUEST;
            case PULL_REQUEST:
                return SnackMessageProto.ProtoMessageType.PULL_REQUEST;
            case COMMIT_OFFSET:
                return SnackMessageProto.ProtoMessageType.COMMIT_OFFSET;
            case FETCH_OFFSET:
                return SnackMessageProto.ProtoMessageType.FETCH_OFFSET;
            case RESPONSE:
                return SnackMessageProto.ProtoMessageType.RESPONSE;
            case HEARTBEAT:
                return SnackMessageProto.ProtoMessageType.HEARTBEAT;
            case DECLARE_EXCHANGE_REQUEST:
                return SnackMessageProto.ProtoMessageType.DECLARE_EXCHANGE_REQUEST;
            case BIND_QUEUE_REQUEST:
                return SnackMessageProto.ProtoMessageType.BIND_QUEUE_REQUEST;
            case UNBIND_QUEUE_REQUEST:
                return SnackMessageProto.ProtoMessageType.UNBIND_QUEUE_REQUEST;
            case DELETE_EXCHANGE_REQUEST:
                return SnackMessageProto.ProtoMessageType.DELETE_EXCHANGE_REQUEST;
            case CREATE_TOPIC_REQUEST:
                return SnackMessageProto.ProtoMessageType.CREATE_TOPIC_REQUEST;
            case DELETE_TOPIC_REQUEST:
                return SnackMessageProto.ProtoMessageType.DELETE_TOPIC_REQUEST;
            case FETCH_TOPIC_METADATA:
                return SnackMessageProto.ProtoMessageType.FETCH_TOPIC_METADATA;
            case AUTHENTICATE_REQUEST:
                return SnackMessageProto.ProtoMessageType.AUTHENTICATE_REQUEST;
            case PREPARE_TRANSACTIONAL_MESSAGE:
                return SnackMessageProto.ProtoMessageType.PREPARE_TRANSACTIONAL_MESSAGE;
            case COMMIT_TRANSACTIONAL_MESSAGE:
                return SnackMessageProto.ProtoMessageType.COMMIT_TRANSACTIONAL_MESSAGE;
            case ROLLBACK_TRANSACTIONAL_MESSAGE:
                return SnackMessageProto.ProtoMessageType.ROLLBACK_TRANSACTIONAL_MESSAGE;
            case CHECK_TRANSACTION_STATE:
                return SnackMessageProto.ProtoMessageType.CHECK_TRANSACTION_STATE;
            case REGISTER_PRODUCER_REQUEST:
                return SnackMessageProto.ProtoMessageType.REGISTER_PRODUCER_REQUEST;
            case GET_CLUSTER_OVERVIEW_REQUEST:
                return SnackMessageProto.ProtoMessageType.GET_CLUSTER_OVERVIEW_REQUEST;
            case LIST_TOPICS_REQUEST:
                return SnackMessageProto.ProtoMessageType.LIST_TOPICS_REQUEST;
            case LIST_CONSUMER_GROUPS_REQUEST:
                return SnackMessageProto.ProtoMessageType.LIST_CONSUMER_GROUPS_REQUEST;
            case LIST_BROKERS_REQUEST:
                return SnackMessageProto.ProtoMessageType.LIST_BROKERS_REQUEST;
            case LIST_USERS_REQUEST:
                return SnackMessageProto.ProtoMessageType.LIST_USERS_REQUEST;
            case LIST_ROLES_REQUEST:
                return SnackMessageProto.ProtoMessageType.LIST_ROLES_REQUEST;
            case CREATE_USER_REQUEST:
                return SnackMessageProto.ProtoMessageType.CREATE_USER_REQUEST;
            case SET_ROLE_PERMISSIONS_REQUEST:
                return SnackMessageProto.ProtoMessageType.SET_ROLE_PERMISSIONS_REQUEST;
            case DELETE_USER_REQUEST:
                return SnackMessageProto.ProtoMessageType.DELETE_USER_REQUEST;
            case DELETE_ROLE_REQUEST:
                return SnackMessageProto.ProtoMessageType.DELETE_ROLE_REQUEST;
            case LIST_EXCHANGES_REQUEST:
                return SnackMessageProto.ProtoMessageType.LIST_EXCHANGES_REQUEST;
            default:
                log.warn("未映射的 MessageType: {}", type);
                return SnackMessageProto.ProtoMessageType.UNRECOGNIZED;
        }
    }

    private static MessageType fromProto(SnackMessageProto.ProtoMessageType type) {
        if (type == null) return null;
        switch (type) {
            case PRODUCE_REQUEST:
                return MessageType.PRODUCE_REQUEST;
            case PULL_REQUEST:
                return MessageType.PULL_REQUEST;
            case COMMIT_OFFSET:
                return MessageType.COMMIT_OFFSET;
            case FETCH_OFFSET:
                return MessageType.FETCH_OFFSET;
            case RESPONSE:
                return MessageType.RESPONSE;
            case HEARTBEAT:
                return MessageType.HEARTBEAT;
            case DECLARE_EXCHANGE_REQUEST:
                return MessageType.DECLARE_EXCHANGE_REQUEST;
            case BIND_QUEUE_REQUEST:
                return MessageType.BIND_QUEUE_REQUEST;
            case UNBIND_QUEUE_REQUEST:
                return MessageType.UNBIND_QUEUE_REQUEST;
            case DELETE_EXCHANGE_REQUEST:
                return MessageType.DELETE_EXCHANGE_REQUEST;
            case CREATE_TOPIC_REQUEST:
                return MessageType.CREATE_TOPIC_REQUEST;
            case DELETE_TOPIC_REQUEST:
                return MessageType.DELETE_TOPIC_REQUEST;
            case FETCH_TOPIC_METADATA:
                return MessageType.FETCH_TOPIC_METADATA;
            case AUTHENTICATE_REQUEST:
                return MessageType.AUTHENTICATE_REQUEST;
            case PREPARE_TRANSACTIONAL_MESSAGE:
                return MessageType.PREPARE_TRANSACTIONAL_MESSAGE;
            case COMMIT_TRANSACTIONAL_MESSAGE:
                return MessageType.COMMIT_TRANSACTIONAL_MESSAGE;
            case ROLLBACK_TRANSACTIONAL_MESSAGE:
                return MessageType.ROLLBACK_TRANSACTIONAL_MESSAGE;
            case CHECK_TRANSACTION_STATE:
                return MessageType.CHECK_TRANSACTION_STATE;
            case REGISTER_PRODUCER_REQUEST:
                return MessageType.REGISTER_PRODUCER_REQUEST;
            case GET_CLUSTER_OVERVIEW_REQUEST:
                return MessageType.GET_CLUSTER_OVERVIEW_REQUEST;
            case LIST_TOPICS_REQUEST:
                return MessageType.LIST_TOPICS_REQUEST;
            case LIST_CONSUMER_GROUPS_REQUEST:
                return MessageType.LIST_CONSUMER_GROUPS_REQUEST;
            case LIST_BROKERS_REQUEST:
                return MessageType.LIST_BROKERS_REQUEST;
            case LIST_USERS_REQUEST:
                return MessageType.LIST_USERS_REQUEST;
            case LIST_ROLES_REQUEST:
                return MessageType.LIST_ROLES_REQUEST;
            case CREATE_USER_REQUEST:
                return MessageType.CREATE_USER_REQUEST;
            case SET_ROLE_PERMISSIONS_REQUEST:
                return MessageType.SET_ROLE_PERMISSIONS_REQUEST;
            case DELETE_USER_REQUEST:
                return MessageType.DELETE_USER_REQUEST;
            case DELETE_ROLE_REQUEST:
                return MessageType.DELETE_ROLE_REQUEST;
            case LIST_EXCHANGES_REQUEST:
                return MessageType.LIST_EXCHANGES_REQUEST;
            case UNRECOGNIZED:
            case UNKNOWN:
            default:
                return null;
        }
    }

    private static SnackMessageProto.ProtoMessageStatus toProto(MessageStatus status) {
        if (status == null) return SnackMessageProto.ProtoMessageStatus.STATUS_UNKNOWN;
        switch (status) {
            case SUCCESS:
                return SnackMessageProto.ProtoMessageStatus.SUCCESS;
            case FAILED:
                return SnackMessageProto.ProtoMessageStatus.FAILED;
            case NOT_FOUND:
                return SnackMessageProto.ProtoMessageStatus.NOT_FOUND;
            case EMPTY:
                return SnackMessageProto.ProtoMessageStatus.EMPTY;
            // 【最终修复】: 添加缺失的 REDIRECT 状态映射
            case REDIRECT:
                return SnackMessageProto.ProtoMessageStatus.REDIRECT;
            default:
                // 对于其他未明确映射的状态，也返回 UNKNOWN，并记录日志
                log.warn("未映射的 MessageStatus: {}", status);
                return SnackMessageProto.ProtoMessageStatus.STATUS_UNKNOWN;
        }
    }

    private static MessageStatus fromProto(SnackMessageProto.ProtoMessageStatus status) {
        if (status == null) return null;
        switch (status) {
            case SUCCESS:
                return MessageStatus.SUCCESS;
            case FAILED:
                return MessageStatus.FAILED;
            case NOT_FOUND:
                return MessageStatus.NOT_FOUND;
            case EMPTY:
                return MessageStatus.EMPTY;
            // 【最终修复】: 添加缺失的 REDIRECT 状态映射
            case REDIRECT:
                return MessageStatus.REDIRECT;
            case UNRECOGNIZED:
            case STATUS_UNKNOWN:
            default:
                return null;
        }
    }

    public static String getPoolStats() {
        return String.format("size=%d, maxSize=%d", poolSize.get(), MAX_POOL_SIZE);
    }
}